package com.jie.component.threadpool.core;

import com.jie.common.util.BaseParameterHolder;
import org.slf4j.MDC;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Callable;

/**
 * @author jie
 * @version 1.0 2025-02-12
 */
public class BaseThreadPool {

    private static final String HOLD_CONTEXT = "holdContext";

    private static final String MDC_CONTEXT = "mdcContext";

    /**
     * 获取父线程的MDC上下文
     * @return
     */
    protected static Map<String, String> getContextForTask() {
        return MDC.getCopyOfContextMap();
    }

    /**
     * 获取父线程的上下文
     * @return
     */
    protected static Map<String,String> getContextForHold() {
        return BaseParameterHolder.getParameterMap();
    }

    /**
     * 装饰线程任务
     * @param runnable
     * @param parentMdcContext
     * @param parentHoldContext
     * @return
     */
    protected static Runnable wrapTask(final Runnable runnable, final Map<String, String> parentMdcContext, final Map<String, String> parentHoldContext) {
        return () -> {
            Map<String, Map<String, String>> preprocess = preprocess(parentMdcContext, parentHoldContext);
            Map<String, String> holdContext = preprocess.get(HOLD_CONTEXT);
            Map<String, String> mdcContext = preprocess.get(MDC_CONTEXT);
            try {
                runnable.run();
            } finally {
                postProcess(mdcContext,holdContext);
            }
        };
    }

    protected static <T> Callable<T> wrapTask(Callable<T> task, final Map<String, String> parentMdcContext, final Map<String, String> parentHoldContext) {
        return () -> {
            Map<String, Map<String, String>> preprocess = preprocess(parentMdcContext, parentHoldContext);
            Map<String, String> holdContext = preprocess.get(HOLD_CONTEXT);
            Map<String, String> mdcContext = preprocess.get(MDC_CONTEXT);
            try {
                return task.call();
            } finally {
                postProcess(mdcContext,holdContext);
            }
        };
    }

    /**
     * 前置处理
     * @param parentMdcContext
     * @param parentHoldContext
     * @return
     */
    private static Map<String, Map<String, String>> preprocess(final Map<String, String> parentMdcContext, final Map<String, String> parentHoldContext) {
        Map<String,Map<String,String>> map = new HashMap<>(8);
        Map<String, String> holdContext = BaseParameterHolder.getParameterMap();
        Map<String, String> mdcContext = MDC.getCopyOfContextMap();
        if (parentMdcContext == null) {
            MDC.clear();
        } else {
            MDC.setContextMap(parentMdcContext);
        }
        if (parentHoldContext == null) {
            BaseParameterHolder.removeParameterMap();
        } else {
            BaseParameterHolder.setParameterMap(parentHoldContext);
        }
        map.put(HOLD_CONTEXT, holdContext);
        map.put(MDC_CONTEXT, mdcContext);
        return map;
    }

    /**
     * 后置处理
     * @param mdcContext
     * @param holdContext
     */
    private static void postProcess(Map<String, String> mdcContext, Map<String, String> holdContext) {
        if (mdcContext == null) {
            MDC.clear();
        } else {
            MDC.setContextMap(mdcContext);
        }

        if (holdContext == null) {
            BaseParameterHolder.removeParameterMap();
        } else {
            BaseParameterHolder.setParameterMap(holdContext);
        }
    }
}
